import torch
import torch.nn as nn
import torch.fft
import numpy as np


class FFTShearletTransform(nn.Module):
    """
    基于 FFT 的 2D 离散 Shearlet 变换，保持输入输出空间分辨率。
    输入: (N, C, H, W)
    输出: 近似系数 (N, C, H, W), 细节系数 (N, C*(num_scales*num_directions-1), H, W)
    """

    def __init__(self, in_channels=32, num_scales=2, num_directions=4, image_size=(128, 160)):
        super(FFTShearletTransform, self).__init__()
        self.in_channels = in_channels
        self.num_scales = num_scales
        self.num_directions = num_directions
        self.H, self.W = image_size

        # 生成频域 Shearlet 滤波器
        self.filters = self._create_shearlet_filters()

    def _create_shearlet_filters(self):
        """
        在频域构造 Shearlet 滤波器，基于伪极坐标网格。
        返回: 频域滤波器列表，形状为 (num_filters, H, W)
        """
        filters = []
        for j in range(self.num_scales):
            scale = 2 ** j  # 尺度因子
            for k in range(self.num_directions):
                # 方向因子（剪切参数）
                shear = k - self.num_directions // 2
                # 构造频域 Shearlet 滤波器（简化为 Meyer 型窗函数）
                filter_freq = self._meyer_shearlet_filter(j, shear)
                filters.append(filter_freq)

        # 转换为 torch 张量并注册为参数
        filter_tensors = [
            torch.tensor(f, dtype=torch.complex64).unsqueeze(0) for f in filters
        ]
        self.filter_bank = nn.ParameterList([
            nn.Parameter(f, requires_grad=False) for f in filter_tensors
        ])
        return self.filter_bank

    def _meyer_shearlet_filter(self, scale_idx, shear):
        """
        构造频域 Meyer 型 Shearlet 滤波器。
        参数:
            scale_idx: 尺度索引
            shear: 剪切参数
        返回: 频域滤波器，形状为 (H, W)
        """
        # 频域网格
        wx = torch.fft.fftfreq(self.H, d=1.0)
        wy = torch.fft.fftfreq(self.W, d=1.0)
        WX, WY = torch.meshgrid(wx, wy, indexing='ij')

        # 尺度参数
        a = 2 ** (-scale_idx)
        # 剪切矩阵作用
        shear_matrix = torch.tensor([[1, -shear], [0, 1]], dtype=torch.float32)
        WX_shear = WX * shear_matrix[0, 0] + WY * shear_matrix[0, 1]
        WY_shear = WX * shear_matrix[1, 0] + WY * shear_matrix[1, 1]

        # Meyer 窗函数（频域）
        radius = torch.sqrt(WX_shear ** 2 + WY_shear ** 2)
        nu = lambda x: torch.where(
            (x >= np.pi / 4) & (x <= np.pi),
            torch.cos(np.pi / 2 * torch.log2(4 * x / np.pi)),
            torch.tensor(0.0)
        )
        window = nu(radius / a)

        # 归一化
        window = window / (window.sum() + 1e-10)
        return window.to(dtype=torch.complex64)

    def forward(self, x):
        """
        应用 FFT 基于的 Shearlet 变换。
        返回: 近似系数 (N, C, H, W), 细节系数 (N, C*(num_filters-1), H, W)
        """
        N, C, H, W = x.shape
        outputs = []

        # 输入转到频域
        x_freq = torch.fft.fft2(x, dim=(-2, -1))

        # 对每个滤波器应用频域卷积
        for filter_idx, filter_freq in enumerate(self.filter_bank):
            if torch.cuda.is_available():
                filter_freq = filter_freq.cuda()
            # 频域乘法（卷积）
            filter_freq = filter_freq.repeat(N, C, 1, 1)
            conv_freq = x_freq * filter_freq
            # 逆 FFT 回到空域
            conv_out = torch.fft.ifft2(conv_freq, dim=(-2, -1)).real
            outputs.append(conv_out)

        # 分离近似系数（第一个滤波器）和细节系数
        approximation = outputs[0]  # 低通滤波器
        details = torch.cat(outputs[1:], dim=1)  # 拼接高通滤波器
        return approximation, details


class FFTShearletPL(nn.Module):
    """
    基于 FFT Shearlet 的处理层，增强高频边缘细节，保持输入输出形状一致。
    参数:
        in_channels: 输入通道数
        num_scales: Shearlet 尺度数量
        num_directions: 每个尺度的方向数量
        image_size: 输入图像尺寸 (H, W)
    """

    def __init__(self, in_channels=32, num_scales=2, num_directions=4, image_size=(128, 160)):
        super(FFTShearletPL, self).__init__()
        self.in_channels = in_channels
        self.num_scales = num_scales
        self.num_directions = num_directions
        self.shearlet = FFTShearletTransform(
            in_channels=in_channels,
            num_scales=num_scales,
            num_directions=num_directions,
            image_size=image_size
        )
        self.softmax = nn.Softmax2d()
        # 调整细节系数的通道数
        self.channel_reducer = nn.Conv2d(
            (num_scales * num_directions - 1) * in_channels,
            in_channels,
            kernel_size=1,
            stride=1,
            padding=0
        )

    def forward(self, input):
        """
        前向传播：应用 FFT Shearlet 变换并增强边缘。
        """
        if torch.cuda.is_available():
            input = input.cuda()

        # 应用 Shearlet 变换
        approximation, details = self.shearlet(input)

        # 调整细节系数通道数
        reduced_details = self.channel_reducer(details)

        # 生成注意力图
        detail_attention = self.softmax(reduced_details)

        # 增强高频细节
        enhanced_details = torch.mul(approximation, detail_attention)

        # 组合近似系数和增强细节
        output = torch.add(approximation, enhanced_details)
        return output


if __name__ == "__main__":
    # 测试模块
    batch_size = 1
    channels = 32
    height, width = 128, 160

    # 创建输入张量
    input_tensor = torch.randn(batch_size, channels, height, width)
    if torch.cuda.is_available():
        input_tensor = input_tensor.cuda()

    # 初始化 FFTShearletPL 模块
    spl = FFTShearletPL(
        in_channels=channels,
        num_scales=2,
        num_directions=4,
        image_size=(height, width)
    )
    if torch.cuda.is_available():
        spl = spl.cuda()

    print(spl)

    # 计算参数量
    num_parameters = sum(p.numel() for p in spl.parameters())
    print(f"参数量: {num_parameters}")

    # 前向传播
    output = spl(input_tensor)

    # 打印输入和输出形状
    print(f"输入形状: {input_tensor.shape}")
    print(f"输出形状: {output.shape}")